# data.py - 数据处理模块
import pandas as pd
import numpy as np
from torch.utils.data import Dataset
from sklearn.preprocessing import StandardScaler


class DataDetime(Dataset):
    """时间序列数据集类"""
    def __init__(self, data, lookback_length, lookforward_length, multi_steps=False):
        self.seq_len = lookback_length
        self.pred_len = lookforward_length
        self.multi_steps = multi_steps
        self.data_y = data
        self.data = data
        print(self.data.shape)

    def __getitem__(self, index):
        s_begin = index
        s_end = s_begin + self.seq_len
        x = self.data[s_begin:s_end]
        if self.multi_steps:
            y = self.data_y[s_end:s_end + self.pred_len, 0]
        else:
            y = self.data_y[s_end + self.pred_len - 1:s_end + self.pred_len, 0]
        return x, y

    def __len__(self):
        return len(self.data) - self.seq_len - self.pred_len + 1


def split_data_cnn(data, train, test, lookback_length):
    """数据分割和预处理函数"""
    # 处理缺失值
    for column in list(data.columns[data.isnull().sum() > 0]):
        data[column].interpolate(method='linear', limit_direction='forward')

    date_column = 'Date' if 'Date' in data.columns else 'date'
    timestamp = data[[date_column]]
    timestamp['date'] = pd.to_datetime(timestamp[date_column])

    cols = list(data.columns)
    cols.remove(date_column)
    data = data[cols].values
    data[:, 0] = np.maximum(data[:, 0], 0)  # 确保第一列非负

    length = len(data)
    num_train = int(length * train)
    num_test = int(length * test)
    num_valid = length - num_test - num_train

    timestamp_train = timestamp[0:num_train]
    timestamp_valid = timestamp[num_train - lookback_length:num_train + num_valid]
    timestamp_test = timestamp[num_train + num_valid - lookback_length:]

    # 标准化处理
    scalar = StandardScaler()
    scalar_y = StandardScaler()
    y = data[0:num_train, 0].reshape(-1, 1)
    scalar_y.fit(y)
    scalar.fit(data[0:num_train])
    data = scalar.transform(data)

    data_train = data[0:num_train]
    data_valid = data[num_train - lookback_length:num_train + num_valid]
    data_test = data[num_train + num_valid - lookback_length:length]

    return data_train, data_valid, data_test, timestamp_train, timestamp_valid, timestamp_test, scalar_y
